{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "15c2d208",
   "metadata": {},
   "source": [
    "# Hijax Variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "99809892",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flax import nnx\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import optax\n",
    "\n",
    "current_mode = nnx.using_hijax()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4220c6f",
   "metadata": {},
   "source": [
    "State propagation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "396a07a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "v = nnx.Variable(jnp.array(0), is_hijax=True)\n",
    "\n",
    "@jax.jit\n",
    "def inc(v):\n",
    "  v[...] += 1\n",
    "\n",
    "print(v[...]); inc(v); print(v[...])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2ab7d801",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:Variable()\u001b[39m. \u001b[34;1mlet\n",
      "    \u001b[39;22mjit[\n",
      "      name=inc\n",
      "      jaxpr={ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:Variable()\u001b[39m. \u001b[34;1mlet\n",
      "          \u001b[39;22mb\u001b[35m:i32[]\u001b[39m = get_variable[avals=(ShapedArray(int32[], weak_type=True),)] a\n",
      "          c\u001b[35m:i32[]\u001b[39m = add b 1:i32[]\n",
      "          _\u001b[35m:i32[]\u001b[39m = get_variable[avals=(ShapedArray(int32[], weak_type=True),)] a\n",
      "          set_variable[\n",
      "            treedef=PyTreeDef(CustomNode(Variable[(('has_ref', False), ('is_hijax', True), ('is_mutable', True))], [*]))\n",
      "            var_type=<class 'flax.nnx.variablelib.Variable'>\n",
      "          ] a c\n",
      "        \u001b[34;1min \u001b[39;22m() }\n",
      "    ] a\n",
      "  \u001b[34;1min \u001b[39;22m() }\n"
     ]
    }
   ],
   "source": [
    "v = nnx.Variable(jnp.array(0), is_hijax=True)\n",
    "print(jax.make_jaxpr(inc)(v))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39070460",
   "metadata": {},
   "source": [
    "Pytree values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fcd0de3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;2;79;201;177mVariable\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 2 (8 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m{\u001b[0m\u001b[38;2;207;144;120m'a'\u001b[0m: Array(0, dtype=int32, weak_type=True), \u001b[38;2;207;144;120m'b'\u001b[0m: Array(2, dtype=int32, weak_type=True)\u001b[38;2;255;213;3m}\u001b[0m,\n",
      "  \u001b[38;2;156;220;254mis_hijax\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n",
      "\u001b[38;2;79;201;177mVariable\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 2 (8 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m{\u001b[0m\u001b[38;2;207;144;120m'a'\u001b[0m: Array(1, dtype=int32, weak_type=True), \u001b[38;2;207;144;120m'b'\u001b[0m: Array(4, dtype=int32, weak_type=True)\u001b[38;2;255;213;3m}\u001b[0m,\n",
      "  \u001b[38;2;156;220;254mis_hijax\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "v = nnx.Variable({'a': jnp.array(0), 'b': jnp.array(2)}, is_hijax=True)\n",
    "\n",
    "@jax.jit\n",
    "def inc_and_double(v):\n",
    "  v['a'] += 1\n",
    "  v['b'] *= 2\n",
    "\n",
    "print(v); inc_and_double(v); print(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0cfe954",
   "metadata": {},
   "source": [
    "Dynamic state structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0d83a130",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before: \u001b[38;2;79;201;177mVariable\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "  \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m{\u001b[0m\u001b[38;2;255;213;3m}\u001b[0m,\n",
      "  \u001b[38;2;156;220;254mis_hijax\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n",
      "After: \u001b[38;2;79;201;177mVariable\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 1 (4 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m{\u001b[0m\u001b[38;2;207;144;120m'y_mean'\u001b[0m: Array(-1.1782329, dtype=float32)\u001b[38;2;255;213;3m}\u001b[0m,\n",
      "  \u001b[38;2;156;220;254mis_hijax\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "rngs = nnx.Rngs(0)\n",
    "x = rngs.uniform((4, 5))\n",
    "w = rngs.normal((5, 3))\n",
    "metrics = nnx.Variable({}, is_hijax=True)\n",
    "\n",
    "@jax.jit\n",
    "def linear(x, w, metrics: nnx.Variable):\n",
    "  y = x @ w\n",
    "  metrics['y_mean'] = jnp.mean(y)\n",
    "  return y\n",
    "\n",
    "print(\"Before:\", metrics)\n",
    "y = linear(x, w, metrics)\n",
    "print(\"After:\", metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0a55df94",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;2;79;201;177mVariable\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 3 (12 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray([1, 2, 3], dtype=int32),\n",
      "  \u001b[38;2;156;220;254mis_hijax\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# set default Variable mode for the rest of the guide\n",
    "nnx.use_hijax(True)\n",
    "\n",
    "variable = nnx.Variable(jnp.array([1, 2, 3]))\n",
    "\n",
    "print(variable)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b2632f1",
   "metadata": {},
   "source": [
    "### Mutability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b7b1f421",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nnx.as_immutable_vars(model) = \u001b[38;2;79;201;177mLinear\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # Param: 3 (12 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mkernel\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;79;201;177mParam\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "    \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;79;201;177mArray\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;156;220;254mshape\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;182;207;169m1\u001b[0m, \u001b[38;2;182;207;169m3\u001b[0m\u001b[38;2;255;213;3m)\u001b[0m, \u001b[38;2;156;220;254mdtype\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mdtype('float32')\u001b[38;2;255;213;3m)\u001b[0m,\n",
      "    \u001b[38;2;156;220;254mis_mutable\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mFalse\u001b[0m,\n",
      "    \u001b[38;2;156;220;254mwas_hijax\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "  \u001b[38;2;255;213;3m)\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n",
      "nnx.as_mutable_vars(model) = \u001b[38;2;79;201;177mLinear\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # Param: 3 (12 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mkernel\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;79;201;177mParam\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "    \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;79;201;177mArray\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;156;220;254mshape\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;182;207;169m1\u001b[0m, \u001b[38;2;182;207;169m3\u001b[0m\u001b[38;2;255;213;3m)\u001b[0m, \u001b[38;2;156;220;254mdtype\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mdtype('float32')\u001b[38;2;255;213;3m)\u001b[0m,\n",
      "    \u001b[38;2;156;220;254mis_hijax\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "  \u001b[38;2;255;213;3m)\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "class Linear(nnx.Module):\n",
    "  def __init__(self, in_features, out_features, rngs: nnx.Rngs):\n",
    "    self.kernel = nnx.Param(rngs.normal((in_features, out_features)))\n",
    "\n",
    "  def __call__(self, x):\n",
    "    return x @ self.kernel\n",
    "\n",
    "model = Linear(1, 3, rngs=nnx.Rngs(0))\n",
    "\n",
    "print(f\"{nnx.as_immutable_vars(model) = !s}\")\n",
    "print(f\"{nnx.as_mutable_vars(model) = !s}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "594cb65e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ImmutableVariableError: Cannot mutate Variable as it is marked as immutable. (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ImmutableVariableError)\n"
     ]
    }
   ],
   "source": [
    "v = nnx.Variable(jnp.array(0))\n",
    "v_immut = nnx.as_immutable_vars(v)\n",
    "assert not v_immut.is_mutable\n",
    "\n",
    "try:\n",
    "  v_immut[...] += 1  # raises an error\n",
    "except Exception as e:\n",
    "  print(f\"{type(e).__name__}: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58692a37",
   "metadata": {},
   "source": [
    "### Ref support"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fcd4fb4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;2;79;201;177mVariable\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 1 (4 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray(0, dtype=int32, weak_type=True),\n",
      "  \u001b[38;2;156;220;254mhas_ref\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n",
      "Ref(0, dtype=int32, weak_type=True)\n"
     ]
    }
   ],
   "source": [
    "v = nnx.Variable(jnp.array(0))\n",
    "v_ref = nnx.as_ref_vars(v)\n",
    "assert v_ref.has_ref\n",
    "print(v_ref)\n",
    "print(v_ref.get_raw_value())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "18256668",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "immutable = \u001b[38;2;79;201;177mVariable\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 1 (4 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray(0, dtype=int32, weak_type=True),\n",
      "  \u001b[38;2;156;220;254mis_mutable\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mFalse\u001b[0m,\n",
      "  \u001b[38;2;156;220;254mhad_ref\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n",
      "mutable = \u001b[38;2;79;201;177mVariable\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 1 (4 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray(0, dtype=int32, weak_type=True),\n",
      "  \u001b[38;2;156;220;254mhas_ref\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "v_immut = nnx.as_immutable_vars(v_ref)\n",
    "assert not v_immut.has_ref\n",
    "print(\"immutable =\", v_immut)\n",
    "\n",
    "v_ref = nnx.as_mutable_vars(v_immut)\n",
    "assert v_ref.has_ref\n",
    "print(\"mutable =\", v_ref)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4e35e75",
   "metadata": {},
   "source": [
    "### Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5400fe58",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nnx.Module):\n",
    "  def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n",
    "    self.linear = Linear(din, dmid, rngs=rngs)\n",
    "    self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n",
    "    self.dropout = nnx.Dropout(0.1, rngs=rngs)\n",
    "    self.linear_out = Linear(dmid, dout, rngs=rngs)\n",
    "\n",
    "  def __call__(self, x):\n",
    "    x = nnx.gelu(self.dropout(self.bn(self.linear(x))))\n",
    "    return self.linear_out(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba980b6b",
   "metadata": {},
   "source": [
    "#### Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "566c4249",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss = 1.000178\n",
      "loss = 0.9700456\n",
      "loss = 0.93967044\n"
     ]
    }
   ],
   "source": [
    "# hijax Variables by default\n",
    "model = Block(2, 64, 3, rngs=nnx.Rngs(0))\n",
    "optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n",
    "\n",
    "@jax.jit\n",
    "def train_step(model, optimizer, x, y):\n",
    "  graphdef, params, nondiff = nnx.split(model, nnx.Param, ...)\n",
    "  def loss_fn(params):\n",
    "    model =  nnx.merge(graphdef, params, nondiff)\n",
    "    return ((model(x) - y) ** 2).mean()\n",
    "\n",
    "  loss, grads = jax.value_and_grad(loss_fn)(nnx.as_immutable_vars(params))  # lojax Variables for jax.grad\n",
    "  optimizer.update(model, grads)\n",
    "\n",
    "  return loss\n",
    "\n",
    "for _ in range(3):\n",
    "  loss = train_step(model, optimizer, x=jnp.ones((10, 2)), y=jnp.ones((10, 3)))\n",
    "  print(f\"{loss = !s}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dea99c1",
   "metadata": {},
   "source": [
    "#### Scan Over Layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d8136be4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: does not work with hijax yet\n",
    "# @jax.vmap\n",
    "# def create_stack(rngs):\n",
    "#   return nnx.as_immutable_vars(Block(2, 64, 2, rngs=rngs))\n",
    "\n",
    "# block_stack = nnx.as_mutable_vars(create_stack(nnx.Rngs(0).fork(split=8)))\n",
    "\n",
    "# def scan_fn(x, block):\n",
    "#   x = block(x)\n",
    "#   return x, None\n",
    "\n",
    "# x = jax.random.uniform(jax.random.key(0), (3, 2))\n",
    "# y, _ = jax.lax.scan(scan_fn, x, block_stack)\n",
    "\n",
    "# print(\"y = \", y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ca18a0d",
   "metadata": {},
   "source": [
    "### Limitations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dd39c79",
   "metadata": {},
   "source": [
    "#### Mutable Outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c6062d19",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: mutable hitypes should use lo_ty_qdd instead\n"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def create_model(rngs):\n",
    "  return Block(2, 64, 3, rngs=rngs)\n",
    "\n",
    "try:\n",
    "  model = create_model(nnx.Rngs(0))\n",
    "except Exception as e:\n",
    "  print(f\"Error:\", e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8bb1e9e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model.linear = \u001b[38;2;79;201;177mLinear\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # Param: 128 (512 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mkernel\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;79;201;177mParam\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
      "    \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;79;201;177mArray\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;156;220;254mshape\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;182;207;169m2\u001b[0m, \u001b[38;2;182;207;169m64\u001b[0m\u001b[38;2;255;213;3m)\u001b[0m, \u001b[38;2;156;220;254mdtype\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mdtype('float32')\u001b[38;2;255;213;3m)\u001b[0m,\n",
      "    \u001b[38;2;156;220;254mis_hijax\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "  \u001b[38;2;255;213;3m)\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def create_model(rngs):\n",
    "  return nnx.as_immutable_vars(Block(2, 64, 3, rngs=rngs))\n",
    "\n",
    "model = nnx.as_mutable_vars(create_model(nnx.Rngs(0)))\n",
    "\n",
    "print(\"model.linear =\", model.linear)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "609bed7c",
   "metadata": {},
   "source": [
    "#### Reference Sharing (aliasing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "045d03c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n"
     ]
    }
   ],
   "source": [
    "# NOTE: doesn't currently fail on the jax side\n",
    "def get_error(f, *args):\n",
    "  try:\n",
    "    return f(*args)\n",
    "  except Exception as e:\n",
    "    return f\"{type(e).__name__}: {e}\"\n",
    "\n",
    "x = nnx.Variable(jnp.array(0))\n",
    "\n",
    "@jax.jit\n",
    "def f(a, b):\n",
    "  ...\n",
    "\n",
    "print(get_error(f, x, x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "bc2e87e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n"
     ]
    }
   ],
   "source": [
    "# NOTE: doesn't currently fail on the jax side\n",
    "class Shared(nnx.Pytree):\n",
    "  def __init__(self):\n",
    "    self.a = nnx.Variable(jnp.array(0))\n",
    "    self.b = self.a\n",
    "    self.c = Linear(1, 1, rngs=nnx.Rngs(0))\n",
    "    self.d = self.c\n",
    "\n",
    "@jax.jit\n",
    "def g(pytree):\n",
    "  ...\n",
    "\n",
    "shared = Shared()\n",
    "\n",
    "print(get_error(g, shared))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6298f3d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Duplicates found:\n",
      "- [('a',), ('b',)]\n",
      "- [('c',), ('d',)]\n"
     ]
    }
   ],
   "source": [
    "print(\"Duplicates found:\")\n",
    "if (all_duplicates := nnx.find_duplicates(shared)):\n",
    "  for duplicates in all_duplicates:\n",
    "    print(\"-\", duplicates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "00854d38",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "before: \u001b[38;2;79;201;177mVariable\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 1 (4 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray(0, dtype=int32, weak_type=True),\n",
      "  \u001b[38;2;156;220;254mis_hijax\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n",
      "after: \u001b[38;2;79;201;177mVariable\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 1 (4 B)\u001b[0m\n",
      "  \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray(10, dtype=int32, weak_type=True),\n",
      "  \u001b[38;2;156;220;254mis_hijax\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n",
      "\u001b[38;2;255;213;3m)\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def h(graphdef, state):\n",
    "  obj = nnx.merge(graphdef, state)\n",
    "  obj.a[...] += 10\n",
    "\n",
    "graphdef, state = nnx.split(shared)\n",
    "print(\"before:\", state.a) # split deduplicates the state\n",
    "\n",
    "h(graphdef, state)\n",
    "\n",
    "print(\"after:\", shared.a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "195296c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# clean up for CI tests\n",
    "_ = nnx.use_hijax(current_mode)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
